package toolkit

import (
	"context"
	"errors"
	"log"
	"mime/multipart"
	"ollama-go/access"
	"ollama-go/config"
	"os"
	"strings"

	"github.com/androidsr/sc-go/model"
	"github.com/gin-gonic/gin"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/tmc/langchaingo/embeddings"
	"github.com/tmc/langchaingo/llms/ollama"
	"github.com/tmc/langchaingo/llms/openai"
	"github.com/tmc/langchaingo/schema"
	"github.com/tmc/langchaingo/vectorstores"
	"github.com/tmc/langchaingo/vectorstores/pgvector"
	_ "gorm.io/driver/postgres"
)

var (
	PG              *PgVector
	collectionTable string
	embeddingTable  string
)

func init() {
	collectionTable = "ollama_collection"
	embeddingTable = "ollama_embedding"
}

type PgVector struct {
	conn      *pgxpool.Pool
	Store     pgvector.Store
	embedding *embeddings.EmbedderImpl
}

func NewOllamaEmbedding() error {
	cfg := config.CFG.Llm.Ollama
	if cfg.EmbeddingDatabase == "" {
		log.Println("否配置数据库，无法启用本地知识库功能")
		return errors.New("否配置数据库，无法启用本地知识库功能")
	}
	PG = new(PgVector)
	var err error
	PG.conn, err = pgxpool.New(context.Background(), cfg.EmbeddingDatabase)
	if err != nil {
		log.Println(err)
		return err
	}
	llm, err := ollama.New(ollama.WithModel(cfg.EmbeddingModel))
	if err != nil {
		return err
	}
	PG.embedding, err = embeddings.NewEmbedder(llm)
	if err != nil {
		log.Println(err)
		return err
	}
	return nil
}

func NewOpenaiEmbedding() error {
	cfg := config.CFG.Llm.Openai
	if cfg.EmbeddingDatabase == "" {
		log.Println("否配置数据库，无法启用本地知识库功能")
		return errors.New("否配置数据库，无法启用本地知识库功能")
	}
	var err error
	PG = new(PgVector)
	PG.conn, err = pgxpool.New(context.Background(), cfg.EmbeddingDatabase)
	if err != nil {
		log.Println(err)
		return err
	}
	llm, err := openai.New(openai.WithBaseURL(cfg.EmbeddingUrl), openai.WithEmbeddingModel(cfg.EmbeddingModel), openai.WithToken(cfg.Token))
	if err != nil {
		log.Println(err)
		return err
	}
	PG.embedding, err = embeddings.NewEmbedder(llm)
	if err != nil {
		log.Println(err)
		return err
	}
	return nil
}

/** 设置STORE参数 **/
func (m *PgVector) AddStore(ctx context.Context, collectionName string, metadata map[string]any) *PgVector {
	if collectionName == "" {
		return nil
	}
	opts := make([]pgvector.Option, 0)
	opts = append(opts,
		pgvector.WithConn(PG.conn),
		pgvector.WithEmbedder(PG.embedding),
		pgvector.WithCollectionTableName(collectionTable),
		pgvector.WithEmbeddingTableName(embeddingTable),
		pgvector.WithCollectionMetadata(metadata),
		pgvector.WithCollectionName(collectionName))
	store, err := pgvector.New(
		ctx,
		opts...,
	)
	if err != nil {
		log.Println(err.Error())
	}
	PG.Store = store
	return m
}

/** 设置STORE参数 **/
func (m *PgVector) SetStore(ctx context.Context, collectionName string) *PgVector {
	if collectionName == "" {
		return nil
	}
	opts := make([]pgvector.Option, 0)
	opts = append(opts,
		pgvector.WithConn(PG.conn),
		pgvector.WithEmbedder(PG.embedding),
		pgvector.WithCollectionTableName(collectionTable),
		pgvector.WithEmbeddingTableName(embeddingTable),
		pgvector.WithCollectionName(collectionName))
	store, err := pgvector.New(
		ctx,
		opts...,
	)
	if err != nil {
		log.Println(err.Error())
	}
	PG.Store = store
	return m
}

/** 增加文档 **/
func (m *PgVector) AddDocument(ctx context.Context, docs []schema.Document) error {
	_, err := m.Store.AddDocuments(ctx, docs)
	return err
}

type SearchDTO struct {
	CollectionName string         `json:"collectionName"`
	Query          string         `json:"query"`
	System         string         `json:"system"`
	Limit          int            `json:"limit"`
	Score          float32        `json:"score"`
	Filter         map[string]any `json:"filter"`
}

/** 查询相似文档 **/
func (m *PgVector) Search(ctx context.Context, query *SearchDTO) ([]schema.Document, error) {
	o := make([]vectorstores.Option, 0)
	if query.Score > 0 {
		o = append(o, vectorstores.WithScoreThreshold(query.Score))
	}
	if query.Filter != nil {
		o = append(o, vectorstores.WithFilters(query.Filter))
	}
	if query.CollectionName != "" {
		o = append(o, vectorstores.WithNameSpace(query.CollectionName))
	}
	docs, err := m.Store.SimilaritySearch(ctx, query.Query, query.Limit, o...)
	return docs, err
}

func (m *PgVector) ReadFormFile(c *gin.Context, cmetadata map[string]any, file *multipart.FileHeader) model.HttpResult {
	collectionName := c.Param("collectionName")
	filename := file.Filename
	os.MkdirAll("/tmp/", 0666)
	dst := "/tmp/" + filename
	c.SaveUploadedFile(file, dst)
	log.Println("保存文件：", dst)
	docs, err := access.GetParsedTextFromUrl(dst, filename[strings.LastIndex(filename, "."):])
	if err != nil {
		return model.NewOK("读取文档内容失败：" + err.Error())
	}
	if len(docs) == 0 {
		return model.NewOK("未读取到任何内容，说明解析不出来")
	}
	err = PG.AddStore(c.Request.Context(), collectionName, cmetadata).AddDocument(c.Request.Context(), docs)
	if err != nil {
		return model.NewFailDefault(err.Error())
	}
	return model.NewOK("处理成功")
}
